import base64

import fastdeploy as fd
import cv2
import os

from config import appSettings


def build_option(device: str, use_trt: bool):
    option = fd.RuntimeOption()
    if device.lower() == "gpu":
        option.use_gpu()
        option.use_paddle_infer_backend()

    if use_trt:
        option.use_trt_backend()
        # If use original Tensorrt, not Paddle-TensorRT,
        # comment the following two lines
        option.enable_paddle_to_trt()
        option.enable_paddle_trt_collect_shape()
        option.set_trt_input_shape("img", [1, 3, 512, 512])

    if device.lower() == "kunlunxin":
        option.use_kunlunxin()
    return option


def generate_filename(path: str, ext: str):
    dirname, filename = os.path.split(path)
    filename_arr = filename.split(".")
    file_name = filename_arr[0]
    return file_name + ext


def base64_encode_img(img):
    with open(img, "rb") as f:
        return base64.b64encode(f.read())


def setup_runtime(model_path):
    model_file = os.path.join(model_path, "model.pdmodel")
    params_file = os.path.join(model_path, "model.pdiparams")
    config_file = os.path.join(model_path, "deploy.yaml")
    model = fd.vision.matting.PPMatting(
        model_file,
        params_file,
        config_file,
        runtime_option=build_option("cpu", False),
    )
    return model


class AIService:
    def __init__(
        self,
        model_path: str = None,
        image_file: str = None,
        save_path: str = appSettings.picture_save_path,
    ):
        self.model_path = model_path
        self.image_file = image_file
        self.save_path = save_path
        self.runtime_option = build_option("cpu", False)
        self.model = setup_runtime(model_path)
        pass

    def predict(self, image_file: str):
        im = cv2.imread(image_file)
        return self.model.predict(im), im

    def remove_bg(self, image_file: str):
        fg_name = generate_filename(image_file, "_fg.png")
        fg_file = self.save_path + "/" + fg_name
        res, im = self.predict(image_file)
        vis_im = fd.vision.vis_matting(im, res, True, 0.5, True)
        cv2.imwrite(fg_file, vis_im)
        # resReturn = {"alpha": base64_encode_img(fg_file)}
        return fg_file
